package com.gitee.wsl.compose.ext.flow

import arrow.optics.Lens
import com.gitee.wsl.flow.state.mapStateFlow
import kotlinx.coroutines.ExperimentalCoroutinesApi
import kotlinx.coroutines.ExperimentalForInheritanceCoroutinesApi
import kotlinx.coroutines.flow.FlowCollector
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.StateFlow
import kotlinx.coroutines.flow.asStateFlow
import kotlinx.coroutines.flow.update


@OptIn(ExperimentalForInheritanceCoroutinesApi::class)
@PublishedApi
internal class TwoWayDerivedStateFlow<T, R>(
    private val upStream: MutableStateFlow<T>,
    private val map: (T) -> R,
    private val unMap: (R) -> T,
) : MutableStateFlow<R> {

    override var value: R
        get() {
            return map(upStream.value)
        }
        set(v) {
            upStream.value = unMap(v)
        }

    override val replayCache: List<R>
        get() = listOf(value)

    private val _sc = MutableStateFlow(0)

    override val subscriptionCount: StateFlow<Int>
        get() = _sc.asStateFlow()

    private val _mappedStream = upStream.mapStateFlow(map)

    override suspend fun collect(collector: FlowCollector<R>): Nothing {
        try {
            _sc.update { it + 1 }
            _mappedStream.collect(collector)
        } finally {
            _sc.update { it - 1 }
        }
    }

    override fun compareAndSet(expect: R, update: R): Boolean {
        return upStream.compareAndSet(
            expect = unMap(expect),
            update = unMap(update),
        )
    }

    @ExperimentalCoroutinesApi
    override fun resetReplayCache() {
        upStream.resetReplayCache()
    }

    override fun tryEmit(value: R): Boolean {
        this.value = value
        return true
    }

    override suspend fun emit(value: R) {
        this.value = value
    }
}

/**
 * NOTE :
 * DON"T USE MutableStateFlow::update
 * If I use the map and unmap does not return equally it will cause to infinite loop
 *
 */
fun <T, R> MutableStateFlow<T>.mapTwoWayStateFlow(
    map: (T) -> R,
    unMap: T.(R) -> T,
): MutableStateFlow<R> {
    return TwoWayDerivedStateFlow(
        upStream = this,
        map = map,
        unMap = {
            unMap(value, it)
        },
    )
}
fun <T, R> MutableStateFlow<T>.mapTwoWayStateFlow(
    lens:Lens<T,R>
): MutableStateFlow<R> {
    return TwoWayDerivedStateFlow(
        upStream = this,
        map = lens::get,
        unMap = {
            lens.set(value,it)
        },
    )
}
